#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
获取文件类型提示词
根据指定的文件类型获取对应的提示词配置
支持先通用识别，再根据识别到的文件类型进行详细处理
"""

import os
import sys
import importlib.util
from typing import Dict, Any, Optional, Tuple

file_type_mappings={
    "通用json":{
        "prompt": """你是一个专业的结构化数据提取助手。请将图片中的关键信息提取为规范的JSON格式，要求如下：
    1. 输出格式必须严格遵循：{'json_result': {'key1': 'value1', 'key2': 'value2', ...}}
    2. 根据图片内容自动识别字段名(key)和对应值(value)，使用最合适的字段名来表示内容
    3. 根据图片类型进行不同处理：
    - 文档/表单：识别各个字段的标签和值，结构化提取
    - 表格：使用表头作为key，单元格内容作为value
    - 卡片/证件：提取各类标识信息
    - 海报/广告：提取关键产品信息、价格、日期等
    4. 所有数值保留原始格式和单位
    5. 对于图片中不存在的信息，相应字段赋值为null
    6. 对于特殊符号和字符，请完整保留
    示例输出：
    {'json_result': {'title': '产品说明书', 'model': 'ABC-123', 'price': '$299.99', 'release_date': '2024-01-15'}}
    现在请处理当前图片内容：
    """
    }
    ,
    "发票": {
        "field_description": """\n发票类型要求：\n1. 价税合计（大写）字段识别"Ⓧ"右边的内容。
        2. 购买方总是在发票的左上方位置，销售方总是在发票的右上方位。
        3. 请识别销售方银行账户的银行名称和账户号码。
        4. 请识别购买方银行账户的银行名称和账户号码。
        5. 请识别发票中出现的"建筑服务"、"旅客运输服务"、"不动产经营租赁服务"、"成品油"和"货物运输服务"等字样，写到specificBusinessTypes字段，若不存在这些字样该字段则为空。
        6. 发票票种只包括专票和普票，写到invoiceType字段。
        7. 请找出备注一栏的内容，填在"remark"字段。\n\n""",
        "fields": {
            "machineNo": {"description": "机器号", "example": "667220920461"},
            "no": {"description": "发票号码", "example": "10043526"},
            "invoiceType": {"description": "发票票种", "example": "专票"},
            "specificBusinessTypes": {"description": "特定业务类型", "example": "建筑服务"},
            "invoiceNumber": {"description": "发票编号", "example": "24122000000005148765"},
            "invoicingDate": {"description": "开票日期", "example": "2015年09月22日"},
            "purchaserName": {"description": "购买方名称", "example": "天津市海纳众诚科技有限公司"},
            "purchaserCode": {"description": "购买方纳税人识别号", "example": "91120116MA074AJDXP"},
            "sellerName": {"description": "销售方名称", "example": "天津中科昊宇科技有限公司"},
            "sellerCode": {"description": "销售方纳税人识别号", "example": "91120116MA07J8594A"},
            "tradeName": {"description": "货物名称（保留\"*\"）", "example": "*信息技术服务*技术服务费"},
            "model": {"description": "规格型号", "example": ""},
            "unit": {"description": "单位", "example": "项"},
            "number": {"description": "数量", "example": ""},
            "unitPrice": {"description": "单价", "example": ""},
            "totalPrice": {"description": "金额（单价×数量）", "example": "95439.62"},
            "taxRate": {"description": "税率", "example": "6%"},
            "taxAmount": {"description": "税额", "example": "5726.38"},
            "totalAmount": {"description": "合计金额", "example": "95439.62"},
            "totalTax": {"description": "合计税额", "example": "5726.38"},
            "payee": {"description": "收款人", "example": "权枢兆基"},
            "reviewer": {"description": "复核", "example": "张金婵"},
            "drawer": {"description": "开票人", "example": "马晓静"},
            "totalBig": {"description": "价税合计(大写)", "example": "壹拾万壹仟壹佰陆拾陆圆整"},
            "totalSmall": {"description": "价税合计(小写)", "example": "101166"},
            "remark": {"description": "备注", "example": ""}
        }
    }
    
}

def get_file_type_mappings():
    return file_type_mappings

def get_prompt_by_type(file_type: str="发票") -> Optional[Dict[str, Any]]:
    """
    根据文件类型从file_type_map_prompt.py中获取相应的提示词配置
    
    Args:
        file_type: 文件类型，如"发票"
    
    Returns:
        包含字段描述和字段信息的字典，如果未找到则返回None
    """
    # 查找指定文件类型的配置
    file_type_mappings=get_file_type_mappings()
    if file_type in file_type_mappings:
        return file_type_mappings[file_type]
    else:
        print(f"未找到文件类型 '{file_type}' 的配置")
        return None

def get_ocr_prompt() -> str:
    """
    获取通用OCR文本提取的提示词
    
    Returns:
        通用OCR文本提取提示词
    """
    ocr_prompt = (
        "请帮我识别这张图片中的全部文字内容，将所有能看到的文字和特殊符号都提取出来，"
        "包括每一行的内容。只需返回纯文本，不要有任何解释和格式化处理。"
        "保持原始的文本布局，每行信息独立成行。"
        "对于数字、金额和日期，请保持原始格式不变。"
        "不要漏掉任何小数点、百分号或货币符号。"
        "如果是表格，请尽量保持表格的行列关系。"
        "不要漏掉任何商品项目，即使有多个商品也要全部识别出来。"
        "务必保留所有特殊符号，如Φ、φ、±、%、㎜、㎝、㎡、¥、$、€、£、₹、₽、®、©、™等，"
        "以及各种科学符号、希腊字母、单位符号等，完全按照原样提取，不要替换为文字描述。"
        "对于商品型号、规格型号中可能包含的特殊字符，请完整保留。"
    )
    return ocr_prompt

def get_two_stage_prompts(file_type: str, ocr_text: str) -> Tuple[str, str]:
    """
    获取两阶段处理的提示词：先OCR识别，再根据文件类型进行结构化处理
    
    Args:
        file_type: 文件类型，如"发票"
        ocr_text: 已识别的OCR文本
        
    Returns:
        (system_prompt, user_prompt)元组
    """
    # 获取指定文件类型的配置
    prompt_info = get_prompt_by_type(file_type)
    
    if not prompt_info:
        raise ValueError(f"无法获取文件类型 '{file_type}' 的提示词信息")
    
    # 提取字段描述和字段信息
    field_description = prompt_info.get('field_description', '')
    fields = prompt_info.get('fields', {})
    
    # 构建系统角色提示词
    system_prompt = (
        "你是一个专业的OCR文本信息提取助手，擅长从OCR文本中提取关键信息，"
        "并将其转换为结构化的JSON数据。请严格按照指定的字段结构提取信息，"
        "确保输出的JSON格式正确且符合要求。对于特殊符号和特殊字符，请完整保留，不要省略或替换。"
    )
    
    # 构建字段列表文本
    fields_info = ""
    for field_name, field_info in fields.items():
        desc = field_info.get('description', '无描述')
        example = field_info.get('example', '无示例')
        fields_info += f"- {field_name}: {desc}, 示例: {example}\n"
    
    # 构建用户提示词，包含OCR文本、字段描述和字段列表
    user_prompt = f"""
请从以下OCR文本中提取信息，并按照指定格式返回JSON：

OCR文本内容:
{ocr_text}

需要提取的字段描述:
{field_description}

字段列表:
{fields_info}

特别注意事项:
1. 请仅返回JSON格式，不要有任何其他文字说明。
2. 确保所有字段值都是字符串格式，即使是数字也要用引号包裹。
3. 如果某个字段在文本中未找到，则该字段值为空字符串。
4. 请完全按照提供的字段列表提取，不要添加或删除任何字段。
5. 当发票中有多个商品项目时，请把所有商品项目的信息都提取出来，不要省略。
   - 对于单位(unit)，如果有多个商品项目，请提取所有商品的单位，用逗号分隔。
   - 对于税率(taxRate)，如果有多个不同税率，请提取所有的税率，用逗号分隔。
   - 对于单价(unitPrice)，如果有多个商品项目，请提取所有商品的单价，用逗号分隔。
   - 对于商品名称(tradeName)，如果有多个商品项目，请提取所有商品名称，用分号分隔。
   - 类似地，对于其他可能有多个值的字段，也请提取所有的值，并用逗号或分号分隔。
6. 对于金额相关数字，请保留原始格式，不要进行四舍五入或变更小数位数。
7. 如果发现金额有误，如税额计算不正确，请按原文提取，不要自行修正。
8. 特别注意：请完整保留所有特殊符号，如Φ、φ、±、%、㎜、㎝、㎡、¥、$、€、£、₹、₽、®、©、™等。
9. 对于包含特殊字符的型号、规格等信息，请保持原样提取，不要进行任何简化或转换。
"""
    
    return system_prompt, user_prompt
